import torch
import pickle
import numpy as np

class LSTMModel(torch.nn.Module):
    def __init__(self, input_size=1, hidden_size=50, num_layers=2):
        super(LSTMModel, self).__init__()
        self.lstm = torch.nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=0.2)
        self.fc = torch.nn.Sequential(
            torch.nn.Linear(hidden_size, 25),
            torch.nn.Linear(25, 1)
        )

    def forward(self, x):
        out, _ = self.lstm(x)
        out = out[:, -1, :]
        out = self.fc(out)
        return out

def load_model_and_scaler(model_path, scaler_path):
    model = LSTMModel()
    model.load_state_dict(torch.load(model_path))
    model.eval()
    with open(scaler_path, 'rb') as f:
        scaler = pickle.load(f)
    return model, scaler

def predict_next_n_hours(model, scaler, input_data, n_hours=24):
    model.eval()
    input_seq = input_data[-24:].copy()
    preds = []
    for _ in range(n_hours):
        scaled_input = scaler.transform(input_seq)
        input_tensor = torch.tensor(scaled_input, dtype=torch.float32).unsqueeze(0)
        with torch.no_grad():
            pred_scaled = model(input_tensor).squeeze().item()
        pred = scaler.inverse_transform(np.array([[pred_scaled]]))[0, 0]
        preds.append(pred)
        input_seq = np.vstack([input_seq[1:], [[pred]]])
    return np.array(preds)
